package net.roseboy.classfinal.util;

import a.a.a.B;
import javassist.ClassPool;
import javassist.CtClass;
import net.roseboy.classfinal.Const;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

public class EncryptUpdateUtil {

    public static void doUpdate(String updateFile,
                                String sourceDir,
                                String targetPath,
                                List<String>  packages,
                                List<String>  excludeClass) {
        Map<String,byte[]> allFileMap = JarUtils.unZipToMap(updateFile, null);
        if(!targetPath.equals("/"))targetPath+="/";
        Map<String,File> updateFiles= new HashMap<>();
        for (String s : allFileMap.keySet()) {
            if(allFileMap.get(s)==null)continue; //跳过目录
            if(s.startsWith(targetPath)){
                File file=new File(sourceDir+File.separator+s.substring(targetPath.length()));
                if(file.exists() && !file.isDirectory()){
                    updateFiles.put(s,file);
                    allFileMap.put(s,IoUtils.readFileToByte(file)) ;
                }
            }
        }

        List<String> classFiles = getClasses(updateFiles.keySet(),packages,excludeClass);
        encryptClass(classFiles,allFileMap);
        clearClassMethod(classFiles,allFileMap);

        packageJar(allFileMap,updateFile);

    }


    /**
     * 压缩成jar
     *
     * @return 打包后的jar绝对路径
     */
    private static void packageJar(Map<String,byte[]> map,String path) {
       IoUtils.writeFile(new File(path),JarUtils.doJar(map));
    }

    /**
     * 清空class文件的方法体，并保留参数信息
     *
     * @param classFiles jar/war 下需要加密的class文件
     */
    private static void clearClassMethod(List<String> classFiles,Map<String,byte[]> map) {
        ClassPool pool = ClassPool.getDefault();
        List<String> classPaths = new ArrayList<>();
        Map<String, CtClass> ctClassMap=new HashMap<>();

        for (Map.Entry<String, byte[]> entry : map.entrySet()) {
            String classFile=entry.getKey();
            byte[] bs=entry.getValue();
            if(bs==null)continue;
            if(classFile.endsWith("class")){
                try {
                    ctClassMap.put(classFile, pool.makeClass(new ByteArrayInputStream(bs)));
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
//        map.forEach((classFile,bs)->{});
        AtomicInteger successCount= new AtomicInteger();
        //[2]修改class方法体，并保存文件
        ClassUtils.errSet.clear();
        for (String classFile : classFiles) {
            String className = resolveClassName(classFile);
            byte[] bts = null;
            try {
                Log.debug("清除方法体: " + className);
                bts = ClassUtils.rewriteAllMethods(pool, ctClassMap.get(classFile));
                successCount.getAndIncrement();
            } catch (Exception e) {
                e.printStackTrace();
                Log.debug("ERROR:" + e.getMessage());
            }
            if (bts != null) {
                map.put(classFile,bts);
            }

        }
        System.out.println("ClassUtils.errSet:"+ClassUtils.errSet);
        System.out.println("清除数量："+successCount.get()+"，总数："+classFiles.size());
    }

    private static List<String> encryptClass(List<String> classFiles,Map<String,byte[]> map) {
        List<String> encryptClasses = new ArrayList<>();
        String newPathRoot="META-INF/"+Const.FILE_NAME+"/";
        //加密另存
        for (String className : classFiles) {
            byte[] bytes = map.get(className);
            if (className.endsWith(".class")) {
                className = resolveClassName(className);
            }
            bytes = B.e(bytes, className.getBytes(StandardCharsets.UTF_8));
            map.put(newPathRoot+className,bytes);
            encryptClasses.add(className);
            Log.debug("加密：" + className);
        }
        return encryptClasses;
    }

    public static List<String> getClasses(Set<String> allFile,
                                          List<String>  packages,
                                          List<String>  excludeClass) {
        List<String> classFiles = new ArrayList<>();
        for (String file : allFile) {
            if (!file.endsWith(".class")) {
                continue;
            }
            //解析出类全名
            String className = resolveClassName(file);
            //判断包名相同和是否排除的类
            if (StrUtils.isMatchs(packages, className, false)
                    && !StrUtils.isMatchs(excludeClass, className, false)) {
                classFiles.add(file);
                Log.debug("待加密: " + file);
            }

        }
        return classFiles;
    }


    /**
     * 根据class的绝对路径解析出class名称或class包所在的路径
     *
     * @param fileName    class绝对路径
     * @return class名称|包所在的路径
     */
    private static String resolveClassName(String fileName) {

        String file = fileName.substring(0, fileName.length() - 6);
        String K_CLASSES = "/classes/";
        String K_LIB = "/lib/";

        String clsName;
        //lib内的的jar包
        if (file.contains(K_LIB)) {
            clsName = file.substring(file.indexOf(Const.LIB_JAR_DIR, file.indexOf(K_LIB))
                    + Const.LIB_JAR_DIR.length() + 1);
        }
        else if (file.contains(K_CLASSES)) {
            clsName = file.substring(file.indexOf(K_CLASSES) + K_CLASSES.length());
        }
        else {
            clsName = file;
        }
        return clsName.replace("/", ".");
    }
}
